#include "dead_code_elimination.h"
#include <map>
using namespace dead_code_elimination;
using namespace cfg;
using namespace live_var;
using namespace std;

ir::IR DeadCodeEliminator::Optimize(ir::IR& ir)
{
    ir::IR new_ir = ir;
    bool changed = true;
    while (changed)
    {
        changed = false;
        for (ir::Function& func: new_ir.funcs)
        {  
            if (OptimizeFunc(func)) { changed = true; }        
        }
    }
    RemoveDeadFunc(new_ir);
    for (ir::Function& func : new_ir.funcs)
    { RemoveDeadAllocInstr(func);}
    return new_ir;
}

bool DeadCodeEliminator::OptimizeFunc(ir::Function& func)
{
    bool changed = false;
    CFG cfg = CFGGenerator().Generate(func);
    FuncLiveness func_liveness = LiveVarAnalyzer().Analyze(cfg);
    if (RemoveDeadInstr(cfg, func_liveness)) { changed = true; }
    if (RemoveDeadBlock(cfg)) { changed = true; }
    if (SimplifyControlFlow(cfg)) { changed = true; }
    func.blocks.clear();
    for (Block& block : cfg.blocks)
    {
        ir::Block ir_block;
        ir_block.label = block.label;
        ir_block.instrs = block.instrs;
        func.blocks.push_back(move(ir_block));
    }
    return changed;
}

void DeadCodeEliminator::RemoveDeadFunc(ir::IR& ir)
{
    map<string, int> func_index_map;
    for (size_t i = 0; i < ir.funcs.size(); ++i)
    {
        ir::Function& func = ir.funcs[i];
        func_index_map[func.name] = i;
    }
    vector<bool> used_funcs(ir.funcs.size());
    vector<int> funcs_index;
    funcs_index.push_back(func_index_map["@main"]);
    while (!funcs_index.empty())
    {
        int func_index = funcs_index.back();
        funcs_index.pop_back();
        if (!used_funcs[func_index])
        {
            used_funcs[func_index] = true;
            for (ir::Block& block : ir.funcs[func_index].blocks)
            {
                for (ir::Instruction& instr : block.instrs)
                {
                    if (instr.op == ir::OP_CALL)
                    {
                        string func_name = get<string>(instr.srcs[0]);
                        funcs_index.push_back(func_index_map[func_name]);
                    }
                }
            }
        }
    }
    for (int i = ir.funcs.size() - 1; i >= 0; --i)
    {
        if (!used_funcs[i]) { ir.funcs.erase(ir.funcs.begin() + i); }
    }
}

void DeadCodeEliminator::RemoveDeadAllocInstr(ir::Function& func)
{
    for (auto alloc_instr = func.alloc_instrs.begin();
        alloc_instr != func.alloc_instrs.end();)
    {
        bool used = false;
        for (ir::Block& block : func.blocks)
        {
            for (ir::Instruction& instr : block.instrs)
            {
                if (instr.op == ir::OP_LOAD)
                {
                    string stack_slot = get<string>(instr.srcs[0]);
                    if (stack_slot == alloc_instr->dest) { used = true; }
                }
                else if (instr.op == ir::OP_STORE)
                {
                    string stack_slot = get<string>(instr.srcs[1]);
                    if (stack_slot == alloc_instr->dest) { used = true; }
                }
                if (used) { break; }
            }
            if (used) { break; }
        }
        if (!used) { alloc_instr = func.alloc_instrs.erase(alloc_instr); }
        else { ++alloc_instr; }
    }
}

bool DeadCodeEliminator::RemoveDeadInstr(
    CFG& cfg, FuncLiveness& func_liveness)
{
    bool changed = false;
    for (Block& block : cfg.blocks)
    {
        BlockLiveness& block_liveness = 
            func_liveness.block_liveness_map[block.label];
        for (int i = block.instrs.size() - 1; i >= 0 ; --i)
        {
            string& var_name = block.instrs[i].dest;
            if (var_name != "")
            {
                InstrLiveness& instr_liveness = 
                    block_liveness.instrs_liveness[i];
                int var_index = func_liveness.var_index_map[var_name];
                if (!instr_liveness.out.test(var_index))
                { 
                    changed = true;
                    block.instrs.erase(block.instrs.begin() + i); 
                }
            }
        }
    }
    return changed;
}

bool DeadCodeEliminator::RemoveDeadBlock(CFG& cfg)
{
    bool changed = false;
    for (int i = cfg.blocks.size() - 1; i >= 0; --i)
    {
        if (cfg.blocks[i].preds.empty() 
            && cfg.blocks[i].label != cfg.entry_block->label)
        {
            changed = true;
            cfg.blocks.erase(cfg.blocks.begin() + i);
        }
    }
    return changed;
}

bool DeadCodeEliminator::SimplifyControlFlow(CFG& cfg)
{
    bool changed = false;
    map<string, Block*> block_map;
    for (Block& block : cfg.blocks)
    { block_map[block.label] = &block; }
    for (Block& block : cfg.blocks)
    {
        ir::Instruction& end_instr = block.instrs.back();
        if (end_instr.op == ir::OP_BR 
            && holds_alternative<int>(end_instr.srcs[0]))
        {
            changed = true;
            int cond = get<int>(end_instr.srcs[0]);
            string then_label = get<string>(end_instr.srcs[1]);
            string else_label = get<string>(end_instr.srcs[2]);
            string taken_label = cond? then_label : else_label;
            ir::Instruction jmp_instr;
            jmp_instr.op = ir::OP_JMP;
            jmp_instr.dest = "";
            jmp_instr.srcs.push_back(taken_label);
            end_instr = move(jmp_instr);
        }
    }
    return changed;
}